import matplotlib.pyplot as plt
import pandas as pd
import numpy as np

rider_num = 18
day_num = 8


def get_money(level):
    total_money = []

    for j in range(1, day_num + 1):
        money = 0
        for i in range(rider_num):
            filename = f"./llm_agent_log/deliver_{i}_record_" + level + ".csv"
            df = pd.read_csv(filename)
            data_frame = pd.read_csv(filename)
            column_money_name = 6
            column_day_name = 11
            filtered_df = df[(df.iloc[:, column_day_name] == j)]
            max_row = filtered_df.iloc[filtered_df.index.argmax()]
            target_value = max_row.iloc[column_money_name]
            # if i == 0 and j == 1:
            #     print(target_value)
            # if(target_value)
            money += target_value

        total_money.append(money)
    # print(total_money)
    # total_money[0] = 0
    for day in range(0, day_num - 1):
        # print(day)
        total_money[day_num - 1 - day] = total_money[day_num - 1 - day] - total_money[day_num - 1 - day - 1]
    total_money.insert(0, 0)
    # print(total_money)
    return total_money


def draw_money_image():
    # 数据处理部分
    # database = SocietyDatabase("AI_Society")
    # regular_money = get_money("regular")
    # imitation_money = get_money("imitation")
    llm_money = get_money("llm")
    #  qlearning_money = get_money("qlearning")
    # qlmemory_money = get_money("qlmemory")
    # mix_money = get_money("mix")

    # plt.plot(regular_money, marker='o', linestyle='-', markersize=5, color='#1E90FF', label='Regular')
    # plt.plot(imitation_money, marker='^', linestyle='-', markersize=5, color='#9400D3', label='Imitation')
    plt.plot(llm_money, marker='^', linestyle='-', markersize=5, color='#32CD32', label='llm')
    # plt.plot(qlearning_money, marker='*', linestyle='-', markersize=5, color='#FFA500', label='RL')
    # plt.plot(qlmemory_money, marker='d', linestyle='-', markersize=5, color='#DC143C', label='RLM')
    # plt.plot(mix_money, marker='x', linestyle='-', markersize=5, color='#32CD32', label='Mix')
    # 显示标签
    plt.legend(loc="lower right")
    plt.title('Riders\' Salary')
    plt.xlabel('Day')
    plt.ylabel('Money')

    plt.savefig("./Storage/llm_result_image/rider_money_line_plot.png", dpi=300)
    # plt.show()


if __name__ == '__main__':
    draw_money_image()
